///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/hf/hf.cpp
///
/// \brief   implementation file for functions related to hartree fock calculations
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include <complex>
#include <cmath>
#include <vector>
#include <iostream>
#include <string>

#include <formic/hf/hf.h>
#include <formic/exception.h>
#include <formic/numeric/numeric.h>
#include <formic/output/output.h>
#include <formic/ham/ham.h>
#include <formic/lapack/interface.h>
#include <formic/mpi/interface.h>
#include <formic/all_data.h>
#include <formic/command/command.h>
#include <formic/diis.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   constructs the one-body reduced density matrix from the orbital coefficient matrix
///
/// \param[in]     n        number of orbitals
/// \param[in]     np       number of particles
/// \param[in]     c        the matrix of orbital coefficients,  size n by n
///
/// \return   the one-body reduced density matrix
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> formic::Matrix<S,int> formic::hf::build_rdm1(const int n,
                                                                const int np,
                                                                const formic::Matrix<S,int> & c) {
  formic::Matrix<S,int> rdm(n, n, formic::zero(S()));
  for (int i = 0; i < np; i++)
  for (int q = 0; q < n; q++)
  for (int p = 0; p < n; p++)
    rdm(p,q) += formic::conj( c(p,i) ) * c(q,i);
  return rdm;
}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   constructs the alpha fock operator
///
/// \param[in]     n        number of orbitals
/// \param[in]     ra       the one body reduced density matrix for alpha, size n*n
/// \param[in]     rb       the one body reduced density matrix for beta , size n*n
/// \param[in]     oei      the one body integrals, size n*n
/// \param[in]     tei      the two body integrals, which are assumed real, in 1122 order, size n*n*n*n
/// \param[out]    f        on exit, the matrix for the fock operator, size n*n
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void formic::hf::build_fock(const int n,
                                                  const std::vector<S> & ra,
                                                  const std::vector<S> & rb,
                                                  const std::vector<S> & oei,
                                                  const std::vector<S> & tei,
                                                        std::vector<S> & f)
{

  // contribution from the one-electron part of the Hamiltonian
  formic::xcopy(n*n, &oei.at(0), 1, &f.at(0), 1);

  // coulomb contribution from alpha
  formic::xgemm('N', 'N', n*n, 1, n*n, formic::unity(S()), &tei.at(0), n*n, &ra.at(0), n*n, formic::unity(S()), &f.at(0), n*n);

  // coulomb contribution from beta
  formic::xgemm('N', 'N', n*n, 1, n*n, formic::unity(S()), &tei.at(0), n*n, &rb.at(0), n*n, formic::unity(S()), &f.at(0), n*n);

  // exchange contribution
  for (int p = 0; p < n; p++)
  for (int q = 0; q < n; q++)
  for (int r = 0; r < n; r++)
  for (int s = 0; s < n; s++)
    f[n*r+q] -= tei[n*n*n*p+n*n*q+n*r+s] * ra[n*p+s];

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   constructs the alpha fock operator using the cholesky decomposed integrals
///
/// \param[in]     n            number of orbitals
/// \param[in]     ra           the one body reduced density matrix for alpha, size n by n
/// \param[in]     rb           the one body reduced density matrix for beta , size n by n
/// \param[in]     ham          the hamiltonian
/// \param[out]    intermed_0   workspace
/// \param[out]    intermed_1   workspace
///
/// \return    the matrix for the fock operator, size n by n
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> formic::Matrix<S,int> formic::hf::build_fock_cholesky(const int n,
                                                                         const formic::Matrix<S,int> & ra,
                                                                         const formic::Matrix<S,int> & rb,
                                                                         const formic::Hamiltonian<S> & ham,
                                                                         std::vector<S> & intermed_0,
                                                                         std::vector<S> & intermed_1) {

  // initialize fock matrix
  formic::Matrix<S,int> f(n, n, formic::zero(S()));

  // get the number of cholesky vectors and size the work vectors accordingly
  const int nchol = ham.n_cholesky();
  if (intermed_0.size() != n*n*nchol) intermed_0.assign(n*n*nchol, formic::zero(S()));
  if (intermed_1.size() != n*n*nchol) intermed_1.assign(n*n*nchol, formic::zero(S()));

  // one-electron contribution
  for (int p = 0; p < n; p++)
  for (int q = 0; q < n; q++)
    f(q,p) = ham.get_oei(p,q);

  if ( nchol > 0 ) {

    // coulomb contribution
    formic::xscal(n*n, formic::zero(S()), &intermed_0.at(0), 1);
    formic::xaxpy(n*n, formic::unity(S()), &ra.at(0,0), 1, &intermed_0.at(0), 1);
    formic::xaxpy(n*n, formic::unity(S()), &rb.at(0,0), 1, &intermed_0.at(0), 1);
    formic::xgemm('T', 'N', nchol, 1, n*n, formic::unity(S()),
                  &ham.cholesky_vecs().at(0), n*n,
                  &intermed_0.at(0), n*n,
                  formic::zero(S()),
                  &intermed_1.at(0), nchol);
    formic::xgemm('N', 'N', n*n, 1, nchol, formic::unity(S()),
                  &ham.cholesky_vecs().at(0), n*n,
                  &intermed_1.at(0), nchol,
                  formic::unity(S()),
                  &f.at(0,0), n*n);

    // exchange contribution
    formic::xgemm('N', 'N', n, n*nchol, n, formic::unity(S()),
                  &ra.at(0,0), n,
                  &ham.cholesky_vecs().at(0), n,
                  formic::zero(S()),
                  &intermed_1.at(0), n);
    for (int u = 0; u < nchol; u++)
    for (int r = 0; r < n; r++)
    for (int p = 0; p < n; p++)
      intermed_0[n*n*u + n*p + r] = intermed_1[n*n*u + n*r + p];
    formic::xgemm('N', 'T', n, n, n*nchol, -formic::unity(S()),
                  &ham.cholesky_vecs().at(0), n,
                  &intermed_0.at(0), n,
                  formic::unity(S()),
                  &f.at(0,0), n);

  }

  // return the fock matrix
  return f;

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   diagonalizes the fock matrix and extracts the new orbital coefficients
///
/// \param[in]     n        number of orbitals
/// \param[in]     np       number of particles
/// \param[in]     f        the matrix for the fock operator, size n by n
/// \param[out]    eval     on exit, the eigenvalues of the the fock matrix, size n
/// \param[out]    c        on exit, the matrix of orbital coefficients, size n by n
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void formic::hf::diagonalize_fock(const int n,
                                                     const int np,
                                                     const formic::Matrix<S,int> & f,
                                                     formic::Vector<double,int> & eval,
                                                     formic::Matrix<S,int> & c) {

  // preparation for diagonalization
  int info;
  const int lwork = 5*n;
  std::vector<S> work(lwork);
  std::vector<double> rwork(3*n);

  // move the fock matrix into the orbital coefficient matrix in preparation for diagonalization
  c = f.clone();

  // take the complex conjugate of the fock matrix.  This way when FORTRAN transposes it the
  // resulting matrix will be equal to the original fock matrix
  formic::xconj(c.size(), &c.at(0,0), 1);

  // diagonalize the fock matrix, after which the eigenvectors will be in the rows of c
  eval.reset(n, 0.0);
  formic::xsyev('V', 'U', n, &c.at(0,0), n, &eval.at(0), &work.at(0), lwork, &rwork.at(0), info);
  if (info != 0)
    throw formic::Exception("formic::xsyev failed with error code %i in formic::hf::diagonalize_fock") % info;

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   constructs the integral arrays used in hartree fock calculations
///
/// \param[in]     ham      object containing the hamiltonian information
/// \param[in]     n        number of orbitals
/// \param[out]    oei_a    on exit, the one body integrals for alpha, size n*n
/// \param[out]    oei_b    on exit, the one body integrals for beta , size n*n
/// \param[out]    tei      on exit, the two body integrals, which are assumed real, in 1122 order, size n*n*n*n
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void formic::hf::build_ints(const formic::Hamiltonian<S> & ham,
                                                  const int n,
                                                  std::vector<S> & oei_a,
                                                  std::vector<S> & oei_b,
                                                  std::vector<S> & tei)
{

//  // quantum chemistry hamiltonian
//  if ( ham.is_qc() ) {
//
//    // one electron integrals
//    for (int p = 0; p < n; p++)
//    for (int q = 0; q < n; q++)
//      oei_a[p*n+q] = oei_b[p*n+q] = ham.get_oei(p,q);
//
//    // two electron integrals
//    for (int p = 0; p < n; p++)
//    for (int q = 0; q < n; q++)
//    for (int r = 0; r < n; r++)
//    for (int s = 0; s < n; s++)
//      tei[ n*n*n*p + n*n*q + n*r + s ] = ham.get_tei(p,q,r,s);
//
//  // hamiltonian built from individual operators
//  } else {
//
//    // zero the integrals
//    pcps::xscal(n*n, formic::zero(S()), &oei_a.at(0), 1);
//    pcps::xscal(n*n, formic::zero(S()), &oei_b.at(0), 1);
//    pcps::xscal(n*n*n*n, formic::zero(S()), &tei.at(0), 1);
//
//    // declare an iterator for looping over the hamiltonian's operators
//    typename std::vector<typename pcps::OpBase<S>::const_ptr>::const_iterator op;
//
//    // extract integrals from the hamiltonian's operators
//    for (op = ham.ops().begin(); op != ham.ops().end(); op++) {
//
//      // raising/lowering operators
//      if (typeid(**op) == typeid(pcps::OpRL<S>)) {
//
//        // check that we do not have any ab contributions
//        if ( (**op).sites().at(0) % 2 != (**op).sites().at(1) % 2 )
//          throw pcps::Exception("RL operators with one alpha and one beta site cannot be used in pcps::hf::build_ints");
//
//        // add the operator's contribution to the one electron integrals
//        const int site1 = (**op).sites().at(0)/2;
//        const int site2 = (**op).sites().at(1)/2;
//        if ((**op).sites().at(0) % 2 == 0 )
//          oei_a[ n * site1 + site2 ] += (**op).coeff();
//        else
//          oei_b[ n * site1 + site2 ] += (**op).coeff();
//
//      // double number operators
//      } else if (typeid(**op) == typeid(pcps::OpNN<S>)) {
//
//        // we currently only support on-site two electron terms like in the Hubbard model
//        const int s0 = (**op).sites().at(0);
//        const int s1 = (**op).sites().at(1);
//        if ( ! ( s0 % 2 == 0 && s1 == s0 + 1 ) &&
//             ! ( s0 % 2 == 1 && s1 == s0 - 1 )    )
//          throw pcps::Exception("NN operators must act on one spatial site in pcps::hf::build_ints");
//
//        // add the operator's contribution to the two electron integrals
//        const int s = (**op).sites().at(0) / 2;
//        tei[ n*n*n*s + n*n*s + n*s + s ] += (**op).coeff();
//
//      // other operators are not currently supported
//      } else {
//        throw pcps::Exception("operator type \"%s\" not supported by pcps::hf::build_ints") % pcps::type_lookup(typeid(**op));
//      }
//
//    }
//
//  }

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   computes the alpha contribution to the hartree fock energy
///
/// \param[in]     n        number of orbitals
/// \param[in]     ham      the hamiltonian
/// \param[in]     f        the matrix for the fock operator, size n by n
/// \param[in]     rdm      the one body reduced density matrix, size n by n
///
/// \return the alpha contribution to the hartree fock energy
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> S formic::hf::alpha_energy(const int n,
                                              const formic::Hamiltonian<S> & ham,
                                              const formic::Matrix<S,int> & f,
                                              const formic::Matrix<S,int> & rdm) {

  // initialize the energy
  S retval = formic::zero(S());

  // compute the energy
  for (int p = 0; p < n; p++)
  for (int q = 0; q < n; q++)
    retval += 0.5 * ( ham.get_oei(p,q) + f(q,p) ) * rdm(q,p);

  // return the energy
  return retval;

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   print the hartree fock orbitals and orbitals energies
///
/// \param[in]     n        number of orbitals
/// \param[in]     na       number of alpha particles
/// \param[in]     evals    the orbital energy eigenvalues, size n
/// \param[in]     c        the matrix of orbital coefficients,  size n by n
/// \param[in]     str      string for the particle type
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void formic::hf::print_orbitals(const int n,
                                                   const int na,
                                                   const formic::Vector<double,int> & evals,
                                                   const formic::Matrix<S,int> & c,
                                                   const std::string & str) {

  if (str.size() > 0)
    formic::of << boost::format("MO coefficients for %s:") % str << std::endl;
  else
    formic::of << boost::format("MO coefficients:") << std::endl;
  formic::of << std::endl;
  for (int i = 0; i < n; i++) {
    formic::of << boost::format("%20.9f  ") % evals[i];
    if (i < na)
      formic::of << boost::format("%4s  ") % "occ";
    else
      formic::of << boost::format("%4s  ") % "vir";
    for (int s = 0; s < n; s++)
      formic::of << boost::format(" %s") % formic::format_number("%12.5f", c.at(s,i));
    formic::of << std::endl;
  }
  formic::of << std::endl;

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   returns a real vector holding the orbital coefficient information
///
/// \param[in]     rhf         whether to perform a restricted or unrestricted calculation
/// \param[in]     n           number of orbitals
/// \param[in]     ca          the alpha orbital coefficient matrix,  size n by n.
/// \param[in]     cb          the beta  orbital coefficient matrix,  size n by n.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> formic::Vector<double,int> formic::hf::get_real_pmt_vec(const bool rhf,
                                                                           const int n,
                                                                           const formic::Matrix<S,int> & ca,
                                                                           const formic::Matrix<S,int> & cb) {

  // get the number of real parameters needed to represent the alpha coefficient matrix
  const int nr = formic::n_real_params(n*n, &ca.at(0,0));

  // get a vector to hold the coefficient matrices' variables in real form
  formic::Vector<double,int> real_vec( ( rhf ? 1 : 2 ) * nr, 0.0);

  // convert alpha variables
  formic::expand_to_real(ca.size(), &ca.at(0,0), &real_vec.at(0));

  // convert beta variables
  if (!rhf)
    formic::expand_to_real(cb.size(), &cb.at(0,0), &real_vec.at(nr));

  // return the real vector of orbital coefficient variables
  return real_vec;

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   gets the alpha and beta coefficients from a real vector of all the coefficients
///
/// \param[in]     rhf         whether to perform a restricted or unrestricted calculation
/// \param[in]     n           number of orbitals
/// \param[in]     rv          the real vector holding all the coefficient variables
/// \param[out]    ca          on exit, the alpha orbital coefficient matrix,  size n by n.
/// \param[out]    cb          on exit, the beta  orbital coefficient matrix,  size n by n.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void formic::hf::from_real_pmt_vec(const bool rhf,
                                                      const int n,
                                                      const formic::Vector<double,int> & rv,
                                                      formic::Matrix<S,int> & ca,
                                                      formic::Matrix<S,int> & cb) {

  // get the number of real parameters needed to represent the alpha coefficient matrix
  const int nr = formic::n_real_params(n*n, (const S *)(0));

  // get alpha variables
  formic::compress_to_complex(n*n, &rv.at(0), &ca.reset(n,n).at(0,0));

  // get beta variables
  if (rhf)
    cb = ca;
  else
    formic::compress_to_complex(n*n, &rv.at(nr), &cb.reset(n,n).at(0,0));

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   computes fock matrix and the hartree fock energy for the supplied orbital coefficients
///
/// \param[in]     rhf         whether to perform a restricted or unrestricted calculation
/// \param[in]     n           number of orbitals
/// \param[in]     na          number of alpha particles
/// \param[in]     nb          number of beta  particles
/// \param[in]     ca          the alpha orbital coefficient matrix
/// \param[in]     cb          the beta  orbital coefficient matrix
/// \param[out]    fa          on exit, the alpha fock matrix
/// \param[out]    fb          on exit, the beta  fock matrix
/// \param[out]    intermed_0   workspace
/// \param[out]    intermed_1   workspace
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> S formic::hf::build_fock_and_energy(const bool rhf,
                                                       const int n,
                                                       const int na,
                                                       const int nb,
                                                       const formic::Hamiltonian<S> & ham,
                                                       const formic::Matrix<S,int> & ca,
                                                       const formic::Matrix<S,int> & cb,
                                                       formic::Matrix<S,int> & fa,
                                                       formic::Matrix<S,int> & fb,
                                                       std::vector<S> & intermed_0,
                                                       std::vector<S> & intermed_1) {

    // build one body density matrices
    const formic::Matrix<S,int> ra = formic::hf::build_rdm1(n, na, ca);
    const formic::Matrix<S,int> rb = ( rhf ? ra : formic::hf::build_rdm1(n, nb, cb) );

    // build fock matrices
    fa = formic::hf::build_fock_cholesky(n, ra, rb, ham, intermed_0, intermed_1);
    fb = ( rhf ? fa : formic::hf::build_fock_cholesky(n, rb, ra, ham, intermed_0, intermed_1) );

    // compute and return energy
    S energy = formic::unity(S()) * ham.e_cor_nuc();
    energy += formic::hf::alpha_energy(n, ham, fa, ra);
    energy += formic::hf::alpha_energy(n, ham, fb, rb);
    return energy;

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   computes the hartree fock energy and orbital coefficients
///
/// \param[in]     rhf         whether to perform a restricted or unrestricted calculation
/// \param[in]     n           number of orbitals
/// \param[in]     na          number of alpha particles
/// \param[in]     nb          number of beta  particles
/// \param[in]     diis_delay  iterations to run before turning on diis
/// \param[in]     ham         object containing the hamiltonian information
/// \param[in]     bias        for each orbital, a bias towards alpha or beta, size n
/// \param[in]     bias_decay  decay rate for the orbital bias
/// \param[in]     thresh      energy convergence threshold
/// \param[in]     maxit       maximum number of self consistent iterations
/// \param[in,out] ca_vec      On entry, the initial guess for the alpha orbital coefficients,  size n*n.
///                            On exit, the optimized matrix of alpha orbital coefficients.
/// \param[in,out] cb_vec      On entry, the initial guess for the beta orbital coefficients,  size n*n.
///                            On exit, the optimized matrix of beta orbital coefficients.
/// \param[out]    energy      On exit, the unrestricted hartree fock energy.
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void formic::hf::run_scf(const bool rhf,
                                            const int n,
                                            const int na,
                                            const int nb,
                                            const int diis_delay,
                                            const formic::Hamiltonian<S> & ham,
                                            const std::vector<double> & bias,
                                            const double bias_decay,
                                            const double thresh,
                                            const int maxit,
                                            std::vector<S> & ca_vec,
                                            std::vector<S> & cb_vec,
                                            S & energy)
{

  // only perform HF on root process
  if (formic::mpi::rank() == 0) {

    // print a greeting
    if (rhf)
      formic::of << boost::format("starting restricted hartree fock iterations") << std::endl;
    else
      formic::of << boost::format("starting unrestricted hartree fock iterations") << std::endl;
    formic::of << std::endl;

    //// allocate and initialize the integrals
    //std::vector<S> oei_a(n*n, formic::zero(S()));
    //std::vector<S> oei_b(n*n, formic::zero(S()));
    //std::vector<S> tei(n*n*n*n, formic::zero(S()));
    //formic::hf::build_ints(ham, n, oei_a, oei_b, tei);

    // check that bias is zero for restricted calculations
    for (int i = 0; rhf && i < bias.size(); i++)
      if ( bias.at(i) != 0.0 )
        throw formic::Exception("bias must be zero for restricted hartree fock");

    // create working copies of the coefficient matrices
    formic::Matrix<S,int> ca, cb;
    ca.reset(n,n) <<= ca_vec;
    cb = ( rhf ? ca : formic::Matrix<S,int>(n,n) <<= cb_vec );

    // variables for the fock matrices
    formic::Matrix<S,int> fa, fb;

    // orbital energy vectors
    formic::Vector<double,int> eval_a(n, 0.0);
    formic::Vector<double,int> eval_b(n, 0.0);

    // prepare some workspace vectors
    std::vector<S> intermed_0, intermed_1;

    // initialize the decay coefficient for the bias
    double decay = 1.0;

    // prepare diis history vectors
    std::vector<formic::Vector<double,int> > pmt_hist;
    std::vector<formic::Vector<double,int> > err_hist;

//    // iterate
//    bool converged = false;
//    for (int iter = 0; iter < maxit; iter++) {
//
//      // check whether diis acceleration has started
//      const bool diis_started = ( pmt_hist.size() > 1 );
//
//      // if diis has started, get diis-improved parameters
//      formic::Vector<double,int> diis_pmt;
//      if (diis_started)
//        diis_pmt = formic::diis_accelerate(pmt_hist, err_hist);
//
//      // save non-improved parameter vector for future use in diis
//      if ( iter >= diis_delay ) {
//        //formic::of << ca.print("%12.6f", "saving coefficient matrix") << std::endl;
//        pmt_hist.push_back( formic::hf::get_real_pmt_vec(rhf, n, ca, cb) );
//      }
//
//      // if diis has started, get diis energy
//      S diis_energy;
//      if (diis_started) {
//        formic::hf::from_real_pmt_vec(rhf, n, diis_pmt, ca, cb);
//        diis_pmt.clear();
//        //formic::of << ca.print("%12.6f", "diis coefficient matrix") << std::endl;
//        diis_energy = formic::hf::build_fock_and_energy(rhf, n, na, nb, ham, ca, cb, fa, fb, intermed_0, intermed_1);
//        formic::hf::from_real_pmt_vec(rhf, n, *pmt_hist.rbegin(), ca, cb);
//      }
//
//      // save previous energy
//      const S last_energy = energy;
//
//      // get fock matrix and energy for current (non-diis-improved) parameters
//      energy = formic::hf::build_fock_and_energy(rhf, n, na, nb, ham, ca, cb, fa, fb, intermed_0, intermed_1);
//
//      // if diis has started, use diis energy
//      if (diis_started)
//        energy = diis_energy;
//
//      // print energy
//      formic::of << boost::format("%4i   %s%s") % iter % formic::format_number("%20.12f", energy) % ( diis_started ? "   diis" : "" ) << std::endl;
//
//      // check for convergence
//      converged = ( iter > 0 && std::abs( energy - last_energy ) < thresh );
//      if (converged) break;
//
//      // add the bias to the fock matrices
//      for (int p = 0; p < n; p++) {
//        fa(p,p) -= formic::unity(S()) * decay * bias.at(p);
//        fb(p,p) += formic::unity(S()) * decay * bias.at(p);
//      }
//
//      // diagonalize fock matrices to find the new orbital coefficients
//      formic::hf::diagonalize_fock(n, na, fa, eval_a, ca);
//      if (rhf) {
//        cb = ca;
//        eval_b = eval_a;
//      } else {
//        formic::hf::diagonalize_fock(n, nb, fb, eval_b, cb);
//      }
//
//      // save parameter difference vector for diis
//      if ( iter >= diis_delay )
//        err_hist.push_back( formic::hf::get_real_pmt_vec(rhf, n, ca, cb).axpy(-1.0, *pmt_hist.rbegin()) );
//
//      // reduce the bias based on the user specified decay rate
//      decay *= bias_decay;
//
//    }

    // iterate
    bool converged = false;
    for (int iter = 0; iter < maxit; iter++) {

      // check whether diis acceleration has started
      const bool diis_started = ( pmt_hist.size() > 1 );

      // if diis has started, get diis-improved parameters
      if (diis_started) {
        formic::hf::from_real_pmt_vec(rhf, n, formic::diis_accelerate(pmt_hist, err_hist), ca, cb);
        const S diis_energy = formic::hf::build_fock_and_energy(rhf, n, na, nb, ham, ca, cb, fa, fb, intermed_0, intermed_1);
        formic::hf::diagonalize_fock(n, na, fa, eval_a, ca);
        if (rhf) {
          cb = ca;
          eval_b = eval_a;
        } else {
          formic::hf::diagonalize_fock(n, nb, fb, eval_b, cb);
        }
      }

      // save parameter vector for future use in diis
      if ( iter >= diis_delay ) {
        //formic::of << ca.print("%12.6f", "saving coefficient matrix") << std::endl;
        pmt_hist.push_back( formic::hf::get_real_pmt_vec(rhf, n, ca, cb) );
      }

      // save previous energy
      const S last_energy = energy;

      // get fock matrix and energy for current (non-diis-improved) parameters
      energy = formic::hf::build_fock_and_energy(rhf, n, na, nb, ham, ca, cb, fa, fb, intermed_0, intermed_1);

      // print energy
      formic::of << boost::format("%4i   %s%s") % iter % formic::format_number("%20.12f", energy) % ( diis_started ? "   diis" : "" ) << std::endl;

      // check for convergence
      converged = ( iter > 0 && std::abs( energy - last_energy ) < thresh );
      if (converged) break;

      // add the bias to the fock matrices
      for (int p = 0; p < n; p++) {
        fa(p,p) -= formic::unity(S()) * decay * bias.at(p);
        fb(p,p) += formic::unity(S()) * decay * bias.at(p);
      }

      // diagonalize fock matrices to find the new orbital coefficients
      formic::hf::diagonalize_fock(n, na, fa, eval_a, ca);
      if (rhf) {
        cb = ca;
        eval_b = eval_a;
      } else {
        formic::hf::diagonalize_fock(n, nb, fb, eval_b, cb);
      }

      // save parameter difference vector for diis
      if ( iter >= diis_delay )
        err_hist.push_back( formic::hf::get_real_pmt_vec(rhf, n, ca, cb).axpy(-1.0, *pmt_hist.rbegin()) );

      // reduce the bias based on the user specified decay rate
      decay *= bias_decay;

    }

    // print results
    formic::of << std::endl;
    if (rhf) {
      if (converged)
        formic::of << boost::format("restricted hartree fock has converged") << std::endl;
      else
        formic::of << boost::format("!!! restricted hartree fock failed to converge !!!") << std::endl;
    } else {
      if (converged)
        formic::of << boost::format("unrestricted hartree fock has converged") << std::endl;
      else
        formic::of << boost::format("!!! unrestricted hartree fock failed to converge !!!") << std::endl;
    }
    formic::of << std::endl;
    formic::of << boost::format("Hartree Fock energy = %s") % formic::format_number("%20.12f", energy) << std::endl;
    formic::of << std::endl;

    // print the mo coefficients and orbital energies
    formic::hf::print_orbitals(n, na, eval_a, ca, "alpha");
    if (!rhf)
      formic::hf::print_orbitals(n, nb, eval_b, cb, "beta");

    // place orbital coefficients in the output vectors
    ca >>= &ca_vec.at(0);
    cb >>= &cb_vec.at(0);

  }

  // reclaim memory
  formic::VectorLock<S,int>::garbage_collect();
  formic::VectorLock<double,int>::garbage_collect();

  // send the results to all processes
  formic::mpi::bcast(energy);
  formic::mpi::bcast(ca_vec);
  formic::mpi::bcast(cb_vec);

}

// explicitly instantiate the function
template void formic::hf::run_scf(const bool rhf,
                                  const int n,
                                  const int na,
                                  const int nb,
                                  const int diis_delay,
                                  const formic::Hamiltonian<double> & ham,
                                  const std::vector<double> & bias,
                                  const double bias_decay,
                                  const double thresh,
                                  const int maxit,
                                  std::vector<double> & ca,
                                  std::vector<double> & cb,
                                  double & energy);
template void formic::hf::run_scf(const bool rhf,
                                  const int n,
                                  const int na,
                                  const int nb,
                                  const int diis_delay,
                                  const formic::Hamiltonian<std::complex<double> > & ham,
                                  const std::vector<double> & bias,
                                  const double bias_decay,
                                  const double thresh,
                                  const int maxit,
                                  std::vector<std::complex<double> > & ca,
                                  std::vector<std::complex<double> > & cb,
                                  std::complex<double> & energy);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "RHF" command, which runs an RHF calculation
///
/// \param[in]     arch         archive to read command arguments from
/// \param[in]     data         data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////

static void formic_rhf_command_function(formic::Archive & arch, formic::AllData & data) {

  // get the name of the wavefunction data to use
  std::string name;
  formic::mpi::read_and_bcast(arch, name, "Failed to read data name for restricted hartree fock orbital coefficients.");

  // get the number of atomic orbitals
  const int nao = data.userinp.get<int>("nao");

  // get the number of alpha and beta particles
  const int na = data.userinp.get<int>("nalpha");
  const int nb = data.userinp.get<int>("nbeta");

  // ensure the number of particles is ok for RHF
  if ( na != nb )
    throw formic::Exception("restricted hartree fock requires an equal number of alpha and beta particles");

  // run the hartree fock calculation
  if (data.userinp.get<std::string>("scalar_type") == "real") {

    // ensure the data vector is the correct size
    if ( data.wd.get_wf(name)->data_vec().size() != nao * nao )
      throw formic::Exception("data vector size incorrect in formic_rhf_command_function.");

    double energy;
    formic::hf::run_scf(true,
                        nao,
                        na,
                        nb,
                        data.userinp.get<int>("hf_diis_delay"),
                        data.hd,
                        data.userinp.get<std::vector<double> >("uhf_bias"),
                        data.userinp.get<double>("uhf_bias_decay"),
                        data.userinp.get<double>("hf_thresh"),
                        data.userinp.get<int>("hf_iter"),
                        data.wd.get_wf(name)->data_vec(),
                        data.wd.get_wf(name)->data_vec(),
                        energy);

  } else if (data.userinp.get<std::string>("scalar_type") == "complex") {

    // ensure the data vector is the correct size
    if ( data.wc.get_wf(name)->data_vec().size() != nao * nao )
      throw formic::Exception("data vector size incorrect in formic_rhf_command_function.");

    std::complex<double> energy;
    formic::hf::run_scf(true,
                        nao,
                        na,
                        nb,
                        data.userinp.get<int>("hf_diis_delay"),
                        data.hc,
                        data.userinp.get<std::vector<double> >("uhf_bias"),
                        data.userinp.get<double>("uhf_bias_decay"),
                        data.userinp.get<double>("hf_thresh"),
                        data.userinp.get<int>("hf_iter"),
                        data.wc.get_wf(name)->data_vec(),
                        data.wc.get_wf(name)->data_vec(),
                        energy);

  } else 
    throw formic::Exception("Unknown scalar type \"%s\" in formic_rhf_command_function.") % data.userinp.get<std::string>("scalar_type");

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "RHF" command, which runs an RHF calculation
///
/// \param[in]     arch         archive to read command arguments from
/// \param[in]     data         data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////

static void formic_uhf_command_function(formic::Archive & arch, formic::AllData & data) {

  // get the name of the wavefunction data to use for the alpha coefficients
  std::string a_name;
  formic::mpi::read_and_bcast(arch, a_name, "Failed to read data name for unrestricted hartree fock alpha orbital coefficients.");

  // get the name of the wavefunction data to use for the beta coefficients
  std::string b_name;
  formic::mpi::read_and_bcast(arch, b_name, "Failed to read data name for unrestricted hartree fock beta orbital coefficients.");

  // get the number of atomic orbitals
  const int nao = data.userinp.get<int>("nao");

  // get the number of alpha and beta particles
  const int na = data.userinp.get<int>("nalpha");
  const int nb = data.userinp.get<int>("nbeta");

  // run the hartree fock calculation
  if (data.userinp.get<std::string>("scalar_type") == "real") {

    // ensure the data vectors are the correct size
    if ( data.wd.get_wf(a_name)->data_vec().size() != nao * nao )
      throw formic::Exception("alpha coefficient matrix size incorrect in formic_uhf_command_function.");
    if ( data.wd.get_wf(b_name)->data_vec().size() != nao * nao )
      throw formic::Exception("beta coefficient matrix size incorrect in formic_uhf_command_function.");

    double energy;
    formic::hf::run_scf(false,
                        nao,
                        na,
                        nb,
                        data.userinp.get<int>("hf_diis_delay"),
                        data.hd,
                        data.userinp.get<std::vector<double> >("uhf_bias"),
                        data.userinp.get<double>("uhf_bias_decay"),
                        data.userinp.get<double>("hf_thresh"),
                        data.userinp.get<int>("hf_iter"),
                        data.wd.get_wf(a_name)->data_vec(),
                        data.wd.get_wf(b_name)->data_vec(),
                        energy);

  } else if (data.userinp.get<std::string>("scalar_type") == "complex") {

    // ensure the data vectors are the correct size
    if ( data.wc.get_wf(a_name)->data_vec().size() != nao * nao )
      throw formic::Exception("alpha coefficient matrix size incorrect in formic_uhf_command_function.");
    if ( data.wc.get_wf(b_name)->data_vec().size() != nao * nao )
      throw formic::Exception("beta coefficient matrix size incorrect in formic_uhf_command_function.");


    std::complex<double> energy;
    formic::hf::run_scf(false,
                        nao,
                        na,
                        nb,
                        data.userinp.get<int>("hf_diis_delay"),
                        data.hc,
                        data.userinp.get<std::vector<double> >("uhf_bias"),
                        data.userinp.get<double>("uhf_bias_decay"),
                        data.userinp.get<double>("hf_thresh"),
                        data.userinp.get<int>("hf_iter"),
                        data.wc.get_wf(a_name)->data_vec(),
                        data.wc.get_wf(b_name)->data_vec(),
                        energy);

  } else 
    throw formic::Exception("Unknown scalar type \"%s\" in formic_uhf_command_function") % data.userinp.get<std::string>("scalar_type");

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Add Hartree Fock commands to the formic command system
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void formic::init_hartree_fock_commands() {

  formic::add_command("RHF", &formic_rhf_command_function);
  formic::add_command("UHF", &formic_uhf_command_function);

}
